package ext;

import ext.trace.geom.Geom;
import ext.trace.kernel.Kernel;
import ext.trace.scene.Hit;
import ext.trace.scene.Scene;
import image.ImageGroup;
import image.ps.PixelShader;
import math.math;
import math.region.Ray3;
import math.region.Surface3;
import static math.vec2.*;
import math.vec2;
import math.vec2i;
import math.vec3;

public class PSRT extends PixelShader
{
    
    public Scene rScene;
    
    public Kernel rKernel;
    
    public vec3 mColorBgnd;
    
    public double mRouletteProb;
    
    public double mDirectProb;
    
    public int mSPP;
    
    public int mTargetI3;
    
    
    public PSRT()
    {
        super();
        this.rScene = null;
        this.rKernel = null;
        this.mColorBgnd = null;
        this.mRouletteProb = -1;
        this.mDirectProb = -1;
        this.mSPP = -1;
        this.mTargetI3 = -1;
    }
    
    
    public PixelShader copy()
    {
        PSRT copy = new PSRT();
        copy.rScene = this.rScene;
        copy.rKernel = this.rKernel;
        copy.mColorBgnd = this.mColorBgnd;
        copy.mRouletteProb = this.mRouletteProb;
        copy.mDirectProb = this.mDirectProb;
        copy.mSPP = this.mSPP;
        copy.mTargetI3 = this.mTargetI3;
        return copy;
    }
    
    
    public void shade(
            ImageGroup aTarget, 
            vec2i aImagePixel, 
            vec2i aImageSize)
    {
        vec2 imageSize = new vec2(aImageSize);
        vec2 imagePixel = new vec2(aImagePixel);
        vec2 rWidth = div(1.0, sub(imageSize, 1.0));
        
        vec3 total = new vec3();
        for (int i = 0; i < this.mSPP; i++)
        {
            vec2 randOffset = this.rRNG.unitDisk();
            randOffset.mul(math.REC_RT2);
            double weight = this.rKernel.eval(randOffset.length(), math.REC_RT2);
            vec2 pixel = add(imagePixel, randOffset);
            vec2 cameraUV = mul(pixel, rWidth);
            Ray3 rayPrimary = this.rScene.rCamera.generate(cameraUV);
            double time = this.rRNG.rand1f(0.0, 1.0);
            vec3 rayContrib = this.rayContrib(rayPrimary, time);
            rayContrib.mul(weight);
            total.add(rayContrib);
        }
        total.div(this.mSPP);
        aTarget.image3(this.mTargetI3).write(aImagePixel, total);
    }
    
    
    private vec3 rayContrib(Ray3 aRay, double time)
    {
        final double unBiasRoulette = 1.0 / (1.0 - this.mRouletteProb);
        if (this.rRNG.rand1f() < this.mRouletteProb)
        {
            return vec3.ZERO.copy();
        }
        Hit hit = this.rScene.castRay(aRay, time);
        if (hit == null)
        {
            return this.mColorBgnd.copy();
        }
        if (hit.rMat.mRadExit > 0.0)
        {
            vec3 emit = hit.rColor.copy();
            emit.mul(hit.rMat.mRadExit);
            emit.mul(hit.rSurface.surfaceArea());
            emit.mul(unBiasRoulette);
            return emit;
        }
        final double indirectProb = 1.0 - this.mDirectProb;
        final double unBiasDirect = 1.0 / this.mDirectProb;
        final double unBiasIndirect = 1.0 / indirectProb;
        if (this.rRNG.rand1f() < this.mDirectProb)
        {
            vec3 direct = this.direct(hit, time);
            direct.mul(unBiasDirect);
            direct.mul(unBiasRoulette);
            return direct;
        }
        else
        {
            vec3 indirect = this.indirect(hit, time);
            indirect.mul(unBiasIndirect);
            indirect.mul(unBiasRoulette);
            return indirect;
        }
    }
    
    private vec3 direct(Hit hit, double time)
    {
        vec3 direct = new vec3();
        for (int i = 0; i < this.rScene.mLights.size(); i++)
        {
            Geom light = this.rScene.mLights.get(i);
            vec2 lightUV = this.rRNG.rand2f();
            vec3 lightContrib = light.mMat.rBRDF.rAlbedo.lookupNN(lightUV);
            lightContrib.mul(light.mMat.mRadExit);
            Surface3 lightSurf = light.lerp(time);
            lightContrib.mul(lightSurf.surfaceArea());
            vec3 L = lightSurf.convertParam_World(lightUV);
            L.sub(hit.rPt);
            L.norm();
            double NL = L.dot(hit.rNorm);
            if (NL < math.EPSILON)
            {
                continue;
            }
            Ray3 occlRay = new Ray3(hit.rPt, L);
            double occl = this.rScene.castOcclusion(occlRay, light, time);
            if (occl < math.EPSILON)
            {
                continue;
            }
            lightContrib.mul(NL * occl / math.PI);
            lightContrib.mul(hit.rColor);
            direct.add(lightContrib);
        }
        return direct;
    }
    
    private vec3 indirect(Hit hit, double time)
    {
        vec3 dirOut = this.rRNG.unitCosHemisphere(hit.rNorm);
        double cos = dirOut.dot(hit.rNorm);
        Ray3 rayIndirect = new Ray3(hit.rPt, dirOut);
        vec3 indirect = this.rayContrib(rayIndirect, time);
        indirect.mul(cos / math.PI);
        indirect.mul(hit.rColor);
        return indirect;
    }
    
}